import json
import time
import re
from datetime import datetime
from typing import Dict, Any
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

# 导入 LangChain 和 LangGraph 记忆与中间件相关的库
from langchain_openai import ChatOpenAI
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain.agents import create_agent, AgentState # 导入 AgentState
from langgraph.checkpoint.memory import InMemorySaver # 内存检查点
from langchain.agents.middleware import before_model # before_model 中间件
from langgraph.runtime import Runtime
from langgraph.graph.message import REMOVE_ALL_MESSAGES
from langchain.messages import RemoveMessage


app = FastAPI()

# 全局变量存储客户端和代理
llm = None
client = None
agent = None
tools = None

# 使用 InMemorySaver 作为内存检查点
# 注意：在生产环境中，请使用数据库支持的 Checkpointer（如 PostgresSaver）
checkpointer = InMemorySaver()

# --- 配置参数 ---
MAX_MESSAGES_TO_KEEP = 10 # 保持的消息总数 (例如: 5轮对话 = 10条消息)

class ChatRequest(BaseModel):
    message: str
    user_id: str = Field(..., description="用于标识会话的唯一用户ID（作为thread_id）")


def parse_tool_args(args_str: str) -> Dict[str, Any]:
    """解析工具调用参数"""
    if not args_str:
        return {}
    try:
        parsed = json.loads(args_str)
        return parsed if isinstance(parsed, dict) else {}
    except:
        try:
            parsed = json.loads(args_str.strip())
            return parsed if isinstance(parsed, dict) else {}
        except:
            # 尝试通过正则表达式解析简单的键值对
            parsed = {}
            try:
                # 匹配 "key": "value"
                for k, v in re.findall(r'"([^"]+)":\s*"([^"]*)"', args_str):
                    parsed[k] = v
                # 匹配 "key": 123 (数字)
                for k, v in re.findall(r'"([^"]+)":\s*(\d+)', args_str):
                    parsed[k] = int(v) if v.isdigit() else v
                # 匹配 key="value" 形式
                if not parsed:
                    for k, v in re.findall(r'(\w+)\s*=\s*"([^"]*)"', args_str):
                        parsed[k] = v
                return parsed
            except:
                return {}

@before_model
def trim_messages(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    """
    短时记忆修剪：只保留最近的 MAX_MESSAGES_TO_KEEP 条消息。
    """
    messages = state["messages"]

    # 如果消息数量没有超过限制，则不做处理
    if len(messages) <= MAX_MESSAGES_TO_KEEP:
        return None

    # 保留最新的消息
    # messages[-MAX_MESSAGES_TO_KEEP:] 表示从后往前取 MAX_MESSAGES_TO_KEEP 条消息
    recent_messages = messages[-MAX_MESSAGES_TO_KEEP:]
    
    # 构建更新状态：先移除所有消息，再添加保留的消息
    # 这样可以确保内存中的历史记录被清除，只留下需要的
    return {
        "messages": [
            RemoveMessage(id=REMOVE_ALL_MESSAGES),
            *recent_messages
        ]
    }


async def initialize_agent():
    """初始化 LLM、客户端和代理"""
    global llm, client, agent, tools
    
    if agent is None:
        llm = ChatOpenAI(
            model="qwen-plus",
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1/",
            api_key="sk-4adf3d4278414c21abf8f26fa9aaa267", # 请替换为您的真实API Key
            temperature=0,
            streaming=True
        )

        client = MultiServerMCPClient(
            {
                "multy_function_server": {
                    "transport": "streamable_http",
                    "url": "http://8.137.115.114:7888/mcp",
                }
            }
        )
        tools = await client.get_tools()
        
        # 使用 create_agent 创建代理，并配置 checkpointer 和中间件
        agent = create_agent(
            llm, 
            tools,
            checkpointer=checkpointer, # 启用短时记忆
            middleware=[trim_messages] # 添加消息修剪中间件
        )

async def stream_chat(request: ChatRequest):
    """流式聊天处理函数"""
    await initialize_agent()

    if not request.user_id:
        raise HTTPException(status_code=400, detail="user_id 字段是必需的。")
    
    # 使用 user_id 作为 thread_id 来配置会话（短时记忆的键）
    config = {"configurable": {"thread_id": request.user_id}}
    
    # 工具调用追踪字典，用于拼接工具调用参数（每个请求独立）
    tool_call_tracker: Dict[str, Dict[str, Any]] = {}
    
    async def generate():
        try:
            async for token, metadata in agent.astream(
                {"messages": [{"role": "user", "content": request.message}]},
                config, # 传入配置，启用短时记忆
                stream_mode="messages",
            ):
                node_name = metadata.get('langgraph_node', '')
                content_blocks = token.content_blocks if hasattr(token, 'content_blocks') else []
                
                # 处理内容块
                for block in content_blocks:
                    block_type = block.get('type')
                    
                    # 处理工具调用块（需要拼接）
                    if block_type == 'tool_call_chunk':
                        call_id = block.get('id', '')
                        tool_name = block.get('name')
                        args_chunk = block.get('args', '')
                        
                        if not call_id:
                            continue
                        
                        # 初始化或更新工具调用追踪
                        if call_id not in tool_call_tracker:
                            tool_call_tracker[call_id] = {
                                'name': tool_name or '',
                                'args': '',
                                'start_time': time.time(),
                                'start_datetime': datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3],
                                'started_reported': False
                            }
                        
                        # 拼接参数和工具名称
                        if args_chunk:
                            tool_call_tracker[call_id]['args'] += args_chunk
                        if tool_name and not tool_call_tracker[call_id]['name']:
                            tool_call_tracker[call_id]['name'] = tool_name
                        
                        # 检测参数是否完整，发送工具调用开始事件
                        if not tool_call_tracker[call_id]['started_reported']:
                            args_str = tool_call_tracker[call_id]['args']
                            has_name = tool_call_tracker[call_id]['name'] or tool_name
                            
                            # 尝试更宽松的 JSON 结束判断
                            # 检查是否有工具名且参数字符串非空，并且包含至少一个完整的 {} 结构
                            if has_name and args_str and (
                                args_str.strip().endswith('}') or 
                                (args_str.count('{') > 0 and args_str.count('}') >= args_str.count('{'))
                            ):
                                parsed_args = parse_tool_args(args_str)
                                yield f"data: {json.dumps({
                                    'type': 'tool_call_start',
                                    'tool_name': tool_call_tracker[call_id]['name'] or tool_name,
                                    'params': parsed_args,
                                    'start_time': tool_call_tracker[call_id]['start_datetime']
                                }, ensure_ascii=False)}\n\n"
                                tool_call_tracker[call_id]['started_reported'] = True
                        
                    # 处理文本内容
                    elif block_type == 'text':
                        # 普通文本内容
                        text = block.get('text', '')
                        if not text:
                            continue
                        
                        # 如果是工具节点返回的文本，说明工具执行完成
                        if node_name == 'tools':
                            # 查找最近完成的工具调用
                            for call_id, call_info in tool_call_tracker.items():
                                if 'end_time' not in call_info and call_info.get('started_reported'):
                                    call_info['end_time'] = time.time()
                                    call_info['end_datetime'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
                                    call_info['duration'] = call_info['end_time'] - call_info['start_time']
                                    
                                    # 解析参数并构建工具调用完成事件
                                    parsed_args = parse_tool_args(call_info['args'])
                                    
                                    yield f"data: {json.dumps({
                                        'type': 'tool_call_end',
                                        'tool_name': call_info['name'],
                                        'params': parsed_args,
                                        'response': {
                                            'content': [{'type': 'text', 'text': text}],
                                            'structuredContent': {'result': text},
                                            'isError': False
                                        },
                                        'start_time': call_info['start_datetime'],
                                        'end_time': call_info['end_datetime'],
                                        'duration_ms': round(call_info['duration'] * 1000, 2)
                                    }, ensure_ascii=False)}\n\n"
                                    break # 假设一个工具执行完成后，会立即返回一个工具消息
                        
                        # 只输出AI回复文本（model节点的文本内容）
                        if node_name == 'model':
                            yield f"data: {json.dumps({
                                'type': 'text',
                                'content': text
                            }, ensure_ascii=False)}\n\n"
                
            # 流结束前，处理所有未完成的工具调用（假设失败或被终端）
            for call_id, call_info in tool_call_tracker.items():
                if 'end_time' not in call_info and call_info.get('started_reported'):
                    call_info['end_time'] = time.time()
                    call_info['end_datetime'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
                    call_info['duration'] = call_info['end_time'] - call_info['start_time']
                    
                    parsed_args = parse_tool_args(call_info['args'])
                    
                    yield f"data: {json.dumps({
                        'type': 'tool_call_end',
                        'tool_name': call_info['name'],
                        'params': parsed_args,
                        'response': {
                            'content': [{'type': 'text', 'text': 'Tool call failed or interrupted.'}],
                            'structuredContent': {'result': None},
                            'isError': True
                        },
                        'start_time': call_info['start_datetime'],
                        'end_time': call_info['end_datetime'],
                        'duration_ms': round(call_info['duration'] * 1000, 2)
                    }, ensure_ascii=False)}\n\n"
            
            # 流结束标记
            yield f"data: {json.dumps({'type': 'done'}, ensure_ascii=False)}\n\n"
            
        except Exception as e:
            # 捕获异常并发送错误事件
            error_message = f"代理执行发生错误: {str(e)}"
            print(f"Error during stream: {error_message}")
            yield f"data: {json.dumps({
                'type': 'error',
                'message': error_message
            }, ensure_ascii=False)}\n\n"
            yield f"data: {json.dumps({'type': 'done'}, ensure_ascii=False)}\n\n"
    
    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"
        }
    )

@app.post("/chat/stream")
async def chat_stream(request: ChatRequest):
    """流式聊天接口"""
    return await stream_chat(request)

@app.get("/health")
async def health():
    """健康检查接口"""
    return {"status": "ok"}

if __name__ == "__main__":
    import uvicorn
    # 异步初始化代理，确保 uvicorn 运行时 agent 已经可用
    import asyncio
    asyncio.run(initialize_agent())
    uvicorn.run(app, host="0.0.0.0", port=8000)